import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F


class TestNet(nn.Module):
    def __init__(self):
        super(TestNet, self).__init__()
        self.op = nn.Embedding(30, 30)

    def forward(self, x):
        out = self.op(x)
        return [out]


if __name__ == "__main__":
    net = TestNet()
    x = np.random.randint(1,5,[3, 3]).astype(np.int32)
    # x = np.ones([3, 3]).astype(np.int32)
    x = torch.IntTensor(x)
    result = net(x)
    print(result[0])
    print(result[0].mean())
    print(result[0].shape)


